import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# 设置matplotlib的字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体为微软雅黑
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题
# 定义SVEIRD模型方程
def sveird_model(t, y, beta, gamma, sigma, delta, alpha, rho):
    S, V, E, I, R, D = y
    N = S + V + E + I + R + D
    dSdt = -beta * S * (I + E) / N - gamma * S
    dVdt = gamma * S - sigma * V * (I + E) / N
    dEdt = beta * S * (I + E) / N + sigma * V * (I + E) / N - delta * E
    dIdt = delta * E - (alpha + rho) * I
    dRdt = alpha * I
    dDdt = rho * I
    return [dSdt, dVdt, dEdt, dIdt, dRdt, dDdt]


# 设定参数
beta = 0.2  # 传染率
gamma = 0.05  # 疫苗接种速率
sigma = 0.1  # 疫苗接种者传染率
delta = 0.1  # 潜伏期逆转率
alpha = 0.1  # 康复率
rho = 0.01  # 病死率

# 设定初始条件
S0, V0, E0, I0, R0, D0 = 990, 10, 0, 1, 0, 0  # 总人口为1000
y0 = [S0, V0, E0, I0, R0, D0]

# 设定时间范围
t_span = (0, 160)  # 分析160天
t_eval = np.linspace(t_span[0], t_span[1], 1000)  # 生成1000个时间点用于评估

# 求解常微分方程组
sol = solve_ivp(sveird_model, t_span, y0, args=(beta, gamma, sigma, delta, alpha, rho), t_eval=t_eval)

# 提取解
S, V, E, I, R, D = sol.y

# 数据可视化
plt.figure(figsize=(12, 8))

plt.plot(t_eval, S / 1000, label='易感者')
plt.plot(t_eval, V / 1000, label='已接种者')
plt.plot(t_eval, E / 1000, label='暴露者')
plt.plot(t_eval, I / 1000, label='感染者')
plt.plot(t_eval, R / 1000, label='康复者')
plt.plot(t_eval, D / 1000, label='死亡者')

plt.xlabel('时间 /天')
plt.ylabel('人数 (千人)')
plt.title('SVEIRD模型可视化结果')
plt.grid(True)
plt.legend()
plt.show()

